package com.simon.study.algorithm.tree;

import com.simon.study.algorithm.tree.TNode.RBNode;
import java.util.ArrayDeque;
import java.util.Deque;
import java.util.Optional;

/**
 * <p>
 *
 * @author simon
 */
public class RBTree {
    RBNode      root;
    int         size;


    public static RBNode find(RBNode node, int data){
        if(node == null){ return null;}

        if(data == node.getData()){ return node; }

        else if(data < node.getData()){
            return find((RBNode) node.getLeft(),  data);
        }else {
            return find((RBNode) node.getRight(), data);
        }
    }


    public static RBNode addNode(RBNode node, int data){
        if(node == null){ return new RBNode(data, RBNode.RED); }
        RBNode target = null;
        if(data <= node.getData()){
            RBNode lchild = (RBNode) node.getLeft();
            target = addNode((RBNode) node.getLeft(), data);
            if(lchild == null){ node.setLeft(target);  target.setParent(node); }
        }else{
            RBNode rchild = (RBNode) node.getRight();
            target = addNode((RBNode) node.getRight(), data);
            if(rchild == null){ node.setRight(target); target.setParent(node); }
        }
        return target;
    }


    public static void main(String[] args) {
        RBTree brtree = new RBTree();

        int[] arr = new int[]{55,66,44,33,11,22,38,46,88,99,30,28,62};

        // brtree.insert(3);
        for (int i = 0; i < arr.length; i++) {
            brtree.insert(arr[i]);
        }

        int[] brr = new int[10];
        for (int i = 0; i < brr.length; i++) {
            brtree.delete(brr[i]);
        }

        brtree.hierachicalPrint();
    }

    public void insert(int data){
        RBNode node = addNode(root, data);
        if(root == null){ root = node; }
        size++;

        fixAfterAdd(node);
    }


    private void fixAfterAdd(RBNode node){
        while (node != root && node.parent().getColor() == RBNode.RED){
            RBNode parent = node.parent(), uncle = node.uncle(), grand = node.grandparent();
            if (uncle != null && uncle.getColor() == RBNode.RED) {
                parent.setColor(RBNode.BLACK);
                uncle.setColor(RBNode.BLACK);
                grand.setColor(RBNode.RED);
                node = grand; continue;
            }

            if(parent == grand.getLeft()){
                // left - right -> left-left
                if(node == parent.getRight()){
                    node = parent;
                    leftRotate(parent);
                }
                grand = node.grandparent();
                grand.setColor(RBNode.RED);
                node.parent().setColor(RBNode.BLACK);
                rightRotate(grand);
            }else{
                if(node == parent.getLeft()){
                    node = parent;
                    rightRotate(parent);
                }
                grand = node.grandparent();
                grand.setColor(RBNode.RED);
                node.parent().setColor(RBNode.BLACK);
                leftRotate(grand);
            }
        }
        root.setColor(RBNode.BLACK);
    }


    public RBNode clear(){
        RBNode node = root;
        root = null; size = 0;
        return node;
    }

    public RBNode successor(RBNode node){
        if(node == null){ return null; }
        else if(node.getRight() != null){
            RBNode s = (RBNode)node.getRight();
            while ( s.getLeft() != null ){
                s = (RBNode) s.getLeft();
            }
            return s;
        }
        else{
            RBNode s = node.parent();
            RBNode c = node;

            while (s != null && c == s.getRight()){
                c = s;
                s = s.parent();
            }
            return s;
        }
    }


    private void leftRotate(RBNode spot){
        RBNode node = (RBNode)AVLTree.leftRotate(spot);
        adjustAfterRotate(spot, node);
        if(root == spot){ root = node; }
    }

    public static void adjustAfterRotate(RBNode spot, RBNode node){
        node.setParent(spot.parent());
        spot.setParent(node);

        Optional.ofNullable(node.parent()).map(x-> x.getLeft() == spot ? x.setLeft(spot) : x.setRight(node));

        Optional.ofNullable(spot.getLeft() ).ifPresent(x -> x.setParent(spot));
        Optional.ofNullable(spot.getRight()).ifPresent(x -> x.setParent(spot));
    }

    private void rightRotate(RBNode spot){
        RBNode node = (RBNode)AVLTree.rightRotate(spot);
        adjustAfterRotate(spot, node);
        if(root == spot){ root = node; }
    }


    public void delete(int data){
        if(size == 0){ return; } RBNode node = find(root, data);

        if(node == null){ return; }
        deleteNode(node); size--;
    }


    public void deleteNode(RBNode node){
        if(node.getLeft() != null && node.getRight() != null){
            RBNode s = successor(node);
            node.setData(s.getData());

            node     = s;
        }

        if (node.parent() == null) {
            root = null; return;
        }

        RBNode target = (RBNode) (node.getLeft() != null ? node.getLeft() : node.getRight());

        RBNode parent = node.parent();
        if(target != null){
            target.setParent(parent);

            if (node == parent.getLeft()) {
                parent.setLeft(target);
            }else{
                parent.setRight(target);
            }

            node.setParent(null).setLeft(null).setRight(null);

            if (node.getColor() == RBNode.BLACK) {
                fixAfterDelete(target);
            }
        }else{
            if (node.getColor() == RBNode.BLACK) {
                fixAfterDelete(node);
            }

            if(node == parent.getLeft()){
                parent.setLeft(null);
            }else{
                parent.setRight(null);
            }
            node.setParent(null);
        }
    }


    public void fixAfterDelete(RBNode node){
        while (node != root && node.getColor() == RBNode.BLACK) {
            RBNode silbing = node.sibling(), parent = node.parent();

            if (node == parent.getLeft()) {
                if (silbing.getColor() == RBNode.RED) {
                    silbing.setColor(RBNode.BLACK);
                    parent.setColor(RBNode.RED);

                    leftRotate(parent);
                    parent  = node.parent();

                    silbing = (RBNode) parent.getRight();
                }

                if(silbing.getLeft() != null || silbing.getRight() != null){
                    if (silbing.getRight() == null) {
                        silbing.setColor(RBNode.RED);
                        ((RBNode)silbing.getLeft()).setColor(RBNode.BLACK);

                        rightRotate(silbing);

                        silbing = (RBNode) parent.getRight();
                    }

                    silbing.setColor(parent.getColor());
                    ((RBNode)silbing.getRight()).setColor(RBNode.BLACK);
                    parent.setColor(RBNode.BLACK);

                    leftRotate(parent);
                    break;
                }else{
                    silbing.setColor(RBNode.RED);
                    node = parent;
                }
            }else{
                if (silbing.getColor() == RBNode.RED) {
                    silbing.setColor(RBNode.BLACK);
                    parent.setColor(RBNode.RED);

                    rightRotate(parent);
                    parent  = node.parent();
                    silbing = (RBNode) parent.getLeft();
                }

                if (silbing.getLeft() != null || silbing.getRight() != null) {
                    if (silbing.getLeft() == null) {
                        silbing.setColor(RBNode.RED);
                        ((RBNode)silbing.getRight()).setColor(RBNode.BLACK);

                        leftRotate(silbing);
                        silbing = (RBNode) parent.getLeft();
                    }

                    silbing.setColor(parent.getColor());
                    parent.setColor(RBNode.BLACK);
                    ((RBNode)silbing.getLeft()).setColor(RBNode.BLACK);
                    rightRotate(parent);
                    break;
                }else{
                    silbing.setColor(RBNode.RED);
                    node = parent;
                }
            }
        }

        node.setColor(RBNode.BLACK);
    }


    public void hierachicalPrint(){
        if(root == null){
            System.out.println("empty tree");
            return;
        }

        Deque<RBNode> stack = new ArrayDeque<>();
        stack.push(root);
        int c = 0;
        while (!stack.isEmpty()){
            RBNode node = stack.peekFirst();

            if(node.getLeft()  != null){ stack.addLast((RBNode) node.getLeft());  c++; }
            if(node.getRight() != null){ stack.addLast((RBNode) node.getRight()); c++; }

            System.out.print(stack.pop() + " ");

            if(stack.size() == c){
                System.out.println();
                c = 0;
            }
        }
    }
}
